import os
import json
from run import run_and_return_test
import pandas as pd
from submission import write_ensemble_submission
import numpy as np
import gc
from config import *

prefix = "/content/drive/MyDrive/lab-result-1/groups"

test_probs = None
config["is_generate"]=False
# submission_path = '/content/drive/MyDrive/ccks2021-wws/submission-Q-s-times.txt'
for similar_fake_num in [50000,100000,200000]:

    for fold_num in ["0","1","2","3","4"]:
        config["group_name"] = f"similar_fake_num={similar_fake_num}"
        config["similar_fake_num"] = similar_fake_num
        # config["steps_per_epoch"] = (similar_fake_num+58000)/config["batch_size"]
        config["train_name"] = config["architecture"].split("/")[-1]+"_fold_"+fold_num

        config["inputs_fold"] = f"/content/drive/MyDrive/ccks2021-wws/lab/data/5fold/fold_"+fold_num


        path = os.path.join(prefix, config["group_name"], config["train_name"])
        config["output_path"] = os.path.join(path, "outputs")
        config_path = os.path.join(path, "config.json")
        if not os.path.exists(path):
            os.makedirs(config["output_path"])
        json.dump(config, open(os.path.join(path, "config.json"), mode="w"))

        if test_probs is None:
            test_probs = run_and_return_test(config_path)
        else:
            test_probs = test_probs+ run_and_return_test(config_path)

        np.save("/content/drive/MyDrive/ccks2021-wws/test_probs-ext.npy",test_probs)
        submission_path = f'/content/drive/MyDrive/ccks2021-wws/submission-similar-{similar_fake_num}.txt'
        write_ensemble_submission(test_probs,submission_path=submission_path)

        gc.collect()

    # Config.fake_label_test_path = submission_path


    
    

